#include <cassert>
#include <cstdlib>
#include <cmath>
#include <ctime>
#include <vector>
#include <iostream>

#include <limits>

#include <cluster.h>

#include <sds/sfile.h>
#include <sds/sdataview.h>

#include "random.h"

using namespace std;
using namespace sds;

typedef SDataView<char> CharView;
typedef SDataView<Point> PointView;

static Random rnd;

int dx[5] = { 0, -1, 0, 0, 1};
int dy[5] = { 0, 0, -1, 1, 0};

int cell = 5;

int state = 6;

double penalty = 48.0;
int length = 150;

double valueProbability(const PointView& points, const CharView& segment,
			int x, int y, char val) {
  double prob = 0.0;
  for (int k = 0; k < cell; ++k) {
    if (x + dx[k] < 0 || x + dx[k] >= segment.getSize()[0]) continue;
    if (y + dy[k] < 0 || y + dy[k] >= segment.getSize()[1]) continue;
    if (segment[x + dx[k]][y + dy[k]] == val) {
      prob += distance_rp(points[x + dx[k]][y + dy[k]], points[x][y]);
      prob -= penalty;
    }
  }
  return prob / (penalty * cell);
}

void Read(const string& filename, PointView& points, vector<SInteger>& size) {
  CharView source[CH];
  SFile file(filename);
  for (int k = 0; k < CH; ++k) {
    source[k] = file.getDataset(k).getDataView<char>();
  }
  
  size = source[0].getSize();
 
  points = PointView(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Point point;
      for (int k = 0; k < CH; ++k) {
	point[k] = source[k][i][j];
      }
      points[i][j] = point;    
    }
  }
}

void Write(const string& filename, CharView& segment) {
  SFile output(filename, SFile::CREATE);
  for (int i = 0; i < segment.getSize()[0]; ++i) {
    for (int j = 0; j < segment.getSize()[1]; ++j) {
      int same = 0;
      for (int k = 1; k < cell; ++k) {
	if (i + dx[k] < 0 || i + dx[k] >= segment.getSize()[0]) continue;
	if (j + dy[k] < 0 || j + dy[k] >= segment.getSize()[1]) continue;
	if (segment[i][j] == segment[i + dx[k]][j + dy[k]]) {
	  ++same;
	}
      }
      if (same == 0) {
	segment[i][j] = 0;
      } else {
	segment[i][j] = segment[i][j] + 1;
      }
    }
  }
  output.createDataset("Segmented Image", UINT8, segment.getSize()).
    setDataView<char>(segment);
}

int main(int argc, const char * argv[]) {
  assert(argc >= 3);
  SFile input(argv[1]);

  length = argc >= 4 ? atoi(argv[3]) : length;
  penalty = argc >= 5 ? atoi(argv[4]) : penalty;

  PointView points;
  vector<SInteger> size;
  
  Read(argv[1], points, size);

  CharView segment(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      segment[i][j] = (int)(state * rnd.uniform()); 
    }
  }
  for (int k = 0; k < length; ++k) {
    if ( (k - 1) * 100 / length < k * 100 / length) 
      cout << k * 100 / length << "%" << endl;
    CharView tmp(size);
    double t = 0.15 / log (1.0 + k * 0.04);
    for (int i = 0; i < size[0]; ++i) {
      for (int j = 0; j < size[1]; ++j) {
	vector<double> probs(state);
	double sum = 0.0;
	for (int l = 0; l < state; ++l) {
	  probs[l] = 
	    exp( -valueProbability(points, segment, i, j, l) / t);
	  sum += probs[l];
	}
	double prob = rnd.uniform() * sum;
	for (int l = 0; l < state; ++l) {
	  if (prob < probs[l]) {
	    tmp[i][j] = l;
	    break;
	  }
	  prob -= probs[l];
	}
      }
    }
    segment.swap(tmp);
  } 
  
  Write(argv[2], segment);

  return 0;
}
